import torch

from models.component.ContextualLoss import ContextualLoss_forward
from models.component.loss import GANLoss
from models.component.vgg19 import VGG19_feature_color_torchversion
from models.networks.basenetwork import BaseNetwork
from models.networks.discriminator import MultiscaleDiscriminator
from models.networks.generator import SPADEGenerator, Generator
from utils import util
from utils.tool import todevice
import torch.nn as nn
import torch.nn.functional as F

class HintModel(nn.Module):
    def __init__(self, opt):
        super(HintModel, self).__init__()
        self.opt = opt
        self.net = torch.nn.ModuleDict(self.initialize_networks())
        self.FloatTensor = torch.cuda.FloatTensor if opt.device == 'cuda' else torch.FloatTensor
        if opt.phase == 'train':
            # init vgg
            self.vggnet_fix = VGG19_feature_color_torchversion(vgg_normal_correct=True)
            self.vggnet_fix.load_state_dict(torch.load('resource/vgg19_conv.pth'))
            self.vggnet_fix.eval()
            for param in self.vggnet_fix.parameters():
                param.requires_grad = False
            self.vggnet_fix.to(opt.device)
            self.perceptual_layer = -1
            # init loss
            self.contextual_forward_loss = ContextualLoss_forward(opt)
            self.MSE_loss = torch.nn.MSELoss()
            self.criterionFeat = torch.nn.L1Loss()
            self.criterionGAN = GANLoss(
                opt.gan_mode, tensor=self.FloatTensor, opt=opt)

    ###---------------------------Init-----------------------------------------------------------------------------##
    def initialize_networks(self):
        net = {}
        G = Generator(self.opt)
        G.initNetwork(self.opt)
        D = MultiscaleDiscriminator(self.opt)
        D.initNetwork(self.opt)
        net['netG'] = G
        net['netD'] = D

        # load model
        if self.opt.phase != 'train' or self.opt.continue_train:
            net['netG'] = util.load_network(net['netG'], 'G', self.opt.which_epoch, self.opt)
            if self.opt.phase == 'train':
                net['netD'] = util.load_network(net['netD'], 'D', self.opt.which_epoch, self.opt)
        return net

    def save(self, epoch):
        util.save_network(self.net['netG'], 'G', epoch, self.opt)
        util.save_network(self.net['netD'], 'D', epoch, self.opt)

    def preprocess_input(self, data):
        sketch = data['sketch']
        image = data['image']
        hints = data['hints']
        sketch_semantic = sketch.clone()
        hint_semantic = hints.clone()
        sketch,image,hints,sketch_semantic,hint_semantic = todevice(self.opt.device,sketch,image,hints,sketch_semantic,hint_semantic)
        return sketch, image, hints, sketch_semantic, hint_semantic

    def inference(self, sketch, hint):
        generate_out = {}
        generate_out['fake_image'] = self.net['netG'](sketch, hint)
        generate_out = {**generate_out}
        return generate_out

    def create_optimizers(self, opt):
        G_params, D_params = list(), list()
        G_params += [{'params': self.net['netG'].parameters(), 'lr': opt.lr * 0.5}]
        if opt.phase == 'train':
            D_params += list(self.net['netD'].parameters())
        if opt.no_TTUR:
            beta1, beta2 = opt.beta1, opt.beta2
            G_lr, D_lr = opt.lr, opt.lr
        else:
            beta1, beta2 = 0, 0.9
            G_lr, D_lr = opt.lr / 2, opt.lr * 2

        optimizer_G = torch.optim.Adam(G_params, lr=G_lr, betas=(beta1, beta2), eps=1e-3)
        optimizer_D = torch.optim.Adam(D_params, lr=D_lr, betas=(beta1, beta2))
        return optimizer_G, optimizer_D
    ###---------------------------Computation----------------------------------------------------------------------##
    def forward(self, data, mode, GforD=None):
        sketch, image, hints, sketch_semantic, hint_semantic = self.preprocess_input(data)
        generated_out = {}
        if mode == 'generator':
            g_loss, generated_out = self.compute_generator_loss(sketch,
                                                                sketch_semantic, image, hints, hint_semantic)
            out = {}
            out['fake_image'] = generated_out['fake_image']
            out['sketch'] = sketch
            out['hints'] = hints
            return g_loss, out
        if mode == 'discriminator':
            d_loss = self.compute_discriminator_loss(sketch_semantic, image, GforD)
            D_losses = {}
            D_losses['D_Fake'] = d_loss['D_Fake']
            D_losses['D_real'] = d_loss['D_real']
            return D_losses
        if mode == 'inference':
            out = {}
            with torch.no_grad():
                out = self.inference(sketch,hints)
                out['sketch'] = sketch
                out['hints'] = hints
            return out
        else:
            raise ValueError("|mode| is invalid")

    ###---------------------------Discriminator--------------------------------------------------------------------##
    def compute_discriminator_loss(self, input_semantics, real_image, GforD):
        D_losses = {}
        with torch.no_grad():
            fake_image = GforD['fake_image'].detach()
            # fake_image.requires_grad_()

        pred_fake, pred_real = self.discriminate(input_semantics, fake_image, real_image)
        D_losses['D_Fake'] = self.criterionGAN(pred_fake, False,
                                               for_discriminator=True) * self.opt.weight_gan
        D_losses['D_real'] = self.criterionGAN(pred_real, True,
                                               for_discriminator=True) * self.opt.weight_gan

        return D_losses

    # Given fake and real image, return the prediction of discriminator
    # for each fake and real image.
    def discriminate(self, input_semantics, fake_image, real_image):
        fake_concat = torch.cat([input_semantics, fake_image], dim=1)
        real_concat = torch.cat([input_semantics, real_image], dim=1)

        # In Batch Normalization, the fake and real images are
        # recommended to be in the same batch to avoid disparate
        # statistics in fake and real images.
        # So both fake and real images are fed to D all at once.
        fake_and_real = torch.cat([fake_concat, real_concat], dim=0)
        discriminator_out = self.net['netD'](fake_and_real)
        pred_fake, pred_real = self.divide_pred(discriminator_out)
        return pred_fake, pred_real

    # Take the prediction of fake and real images from the combined batch
    def divide_pred(self, pred):
        # the prediction contains the intermediate outputs of multiscale GAN,
        # so it's usually a list
        if type(pred) == list:
            fake = []
            real = []
            for p in pred:
                fake.append([tensor[:tensor.size(0) // 2] for tensor in p])
                real.append([tensor[tensor.size(0) // 2:] for tensor in p])
        else:
            fake = pred[:pred.size(0) // 2]
            real = pred[pred.size(0) // 2:]

        return fake, real

    ###---------------------------Generator--------------------------------------------------------------------##
    def generate_fake(self, sketch_semantic, real_image, hint_semantic):
        generate_out = {}
        ref_relu1_1, ref_relu2_1, ref_relu3_1, ref_relu4_1, ref_relu5_1 = self.vggnet_fix(sketch_semantic,
                                                                                          ['r12', 'r22', 'r32', 'r42',
                                                                                           'r52'], preprocess=True)
        generate_out['hint_features'] = [ref_relu1_1, ref_relu2_1, ref_relu3_1, ref_relu4_1, ref_relu5_1]
        generate_out['real_features'] = self.vggnet_fix(real_image, ['r12', 'r22', 'r32', 'r42', 'r52'],
                                                        preprocess=True)

        generate_out['fake_image'] = self.net['netG'](sketch_semantic, hint_semantic)
        generate_out = {**generate_out}
        return generate_out

    def get_ctx_loss(self, source, target):
        contextual_style5_1 = torch.mean(self.contextual_forward_loss(source[-1], target[-1].detach())) * 8
        contextual_style4_1 = torch.mean(self.contextual_forward_loss(source[-2], target[-2].detach())) * 4
        contextual_style3_1 = torch.mean(self.contextual_forward_loss(F.avg_pool2d(source[-3], 2), F.avg_pool2d(target[-3].detach(), 2))) * 2
        if self.opt.use_22ctx:
            contextual_style2_1 = torch.mean(self.contextual_forward_loss(F.avg_pool2d(source[-4], 4), F.avg_pool2d(target[-4].detach(), 4))) * 1
            return contextual_style5_1 + contextual_style4_1 + contextual_style3_1 + contextual_style2_1
        return contextual_style5_1 + contextual_style4_1 + contextual_style3_1

    def compute_generator_loss(self, sketch, sketch_semantic, image, hints, hint_semantic):
        G_losses = {}
        generate_out = self.generate_fake(sketch_semantic, image, hint_semantic)

        # adv loss
        pred_fake, pred_real = self.discriminate(sketch_semantic, generate_out['fake_image'], image)
        G_losses['GAN'] = self.criterionGAN(pred_fake, True, for_discriminator=False) * self.opt.weight_gan

        # feat loss for D
        num_D = len(pred_fake)
        GAN_Feat_loss = self.FloatTensor(1).fill_(0)
        for i in range(num_D):  # for each discriminator
            # last output is the final prediction, so we exclude it
            num_intermediate_outputs = len(pred_fake[i]) - 1
            for j in range(num_intermediate_outputs):  # for each layer output
                unweighted_loss = self.criterionFeat(
                    pred_fake[i][j], pred_real[i][j].detach())
                GAN_Feat_loss += unweighted_loss * self.opt.lambda_feat / num_D
        G_losses['GAN_Feat_l1'] = GAN_Feat_loss

        # l1 loss for perceptual
        fake_features = self.vggnet_fix(generate_out['fake_image'], ['r12', 'r22', 'r32', 'r42', 'r52'],
                                        preprocess=True)
        sample_weights = (image[:, 0, 0, 0] / (sum(image[:, 0, 0, 0]) + 1e-5)).unsqueeze(-1).unsqueeze(
            -1).unsqueeze(-1)
        weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
        loss = 0
        for i in range(len(generate_out['real_features'])):
            loss += weights[i] * util.weighted_l1_loss(fake_features[i], generate_out['real_features'][i].detach(),
                                                       sample_weights)
        G_losses['perc'] = loss * self.opt.lambda_vgg * self.opt.fm_ratio

        # mse loss
        feat_loss = util.mse_loss(fake_features[self.perceptual_layer],
                                  generate_out['real_features'][self.perceptual_layer].detach())
        G_losses['mse'] = feat_loss * self.opt.weight_mse

        # contextual loss
        G_losses['contextual'] = self.get_ctx_loss(fake_features,
                                                   generate_out['real_features']) * self.opt.lambda_vgg * self.opt.ctx_w
        return G_losses, generate_out